import io
import re
import os
import glob
import random
import base64
import threading

from typing import Optional
from PIL import Image
from datetime import datetime
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import JSONResponse

from .utils.adb import ADBManager

# 创建路由
router = APIRouter(prefix="/data", tags=["data collection"])

# 数据采集配置
DATA_ROOT = r"files\hand\data"
SERVICE_DATA_PATH = os.path.join(DATA_ROOT, "ServiceDataLoggerRecord")

# 数据采集状态
collection_status = {
    "is_collecting": False,
    "last_collection_time": None,
    "last_collection_count": 0,
    "current_collection_thread": None
}


def get_latest_collection_session():
    """获取最新的数据采集会话目录"""
    if not os.path.exists(SERVICE_DATA_PATH):
        return None

    sessions = glob.glob(os.path.join(SERVICE_DATA_PATH, "yvrdatalogger*"))
    if not sessions:
        return None

    # 按创建时间排序，获取最新的会话
    sessions.sort(key=os.path.getctime, reverse=True)
    return sessions[0]


def count_images_in_session(session_path):
    """统计会话中的图片数量"""
    total_count = 0
    camera_folders = glob.glob(os.path.join(session_path, "Camera*"))

    for camera_folder in camera_folders:
        images_path = os.path.join(camera_folder, "images", "*.bmp")
        images = glob.glob(images_path)
        total_count += len(images)

    return total_count


def get_all_collection_sessions():
    """获取所有数据采集会话"""
    if not os.path.exists(SERVICE_DATA_PATH):
        return []

    sessions = glob.glob(os.path.join(SERVICE_DATA_PATH, "yvrdatalogger*"))
    sessions_info = []

    for session_path in sessions:
        session_name = os.path.basename(session_path)
        # 从会话名称中提取时间戳
        timestamp_match = re.search(r'yvrdatalogger(\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2})', session_name)
        if timestamp_match:
            timestamp = timestamp_match.group(1)
            image_count = count_images_in_session(session_path)
            creation_time = os.path.getctime(session_path)

            sessions_info.append({
                "timestamp": timestamp,
                "session_name": session_name,
                "image_count": image_count,
                "creation_time": creation_time,
                "path": session_path
            })

    # 按创建时间排序
    sessions_info.sort(key=lambda x: x["creation_time"], reverse=True)
    return sessions_info


def get_session_by_timestamp(timestamp):
    """根据时间戳获取会话"""
    if not timestamp:
        return get_latest_collection_session()

    session_pattern = f"yvrdatalogger{timestamp}"
    session_path = os.path.join(SERVICE_DATA_PATH, session_pattern)

    if os.path.exists(session_path):
        return session_path
    else:
        # 尝试查找包含该时间戳的会话
        sessions = glob.glob(os.path.join(SERVICE_DATA_PATH, f"*{timestamp}*"))
        if sessions:
            return sessions[0]
        else:
            return None


def get_images_by_id(id_=None):
    """根据时间戳获取前6个Camera的图片"""
    session_path = get_latest_collection_session()

    if not session_path:
        return None

    # 获取所有Camera文件夹
    camera_folders = glob.glob(os.path.join(session_path, "Camera*"))
    camera_folders.sort()  # 按Camera0, Camera1, ... 排序

    images_data = []

    for camera_folder in camera_folders[:6]:  # 只取前6个Camera
        camera_name = os.path.basename(camera_folder)
        images_path = os.path.join(camera_folder, "images", "*.bmp")
        image_files = glob.glob(images_path)

        if image_files:
            # 按文件名（时间戳）排序，取最早的一张
            image_files.sort()
            if isinstance(id_, int) and 0 <= id_ < len(image_files):
                earliest_image = image_files[id_]
            else:
                id_ = random.randint(0, len(image_files) - 1)
                earliest_image = image_files[id_]
            # 读取图片并转换为base64
            try:
                with Image.open(earliest_image) as img:
                    # 转换为RGB模式（如果是灰度图）
                    if img.mode != 'RGB':
                        img = img.convert('RGB')

                    # 调整图片大小（可选，避免返回过大图片）
                    max_size = (320, 240)
                    img.thumbnail(max_size, Image.Resampling.LANCZOS)

                    buffer = io.BytesIO()
                    img.save(buffer, format="JPEG", quality=85)
                    img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')

                    # 获取图片时间戳（文件名）
                    image_timestamp = os.path.basename(earliest_image).replace('.bmp', '')

                    images_data.append({
                        "camera": camera_name,
                        "image_base64": img_base64,
                        "timestamp": image_timestamp,
                        "format": "JPEG"  # 转换为JPEG以减少大小
                    })
            except Exception as e:
                print(f"Error processing image {earliest_image}: {e}")
                continue

    return {
        "session_timestamp": os.path.basename(session_path).replace('yvrdatalogger', ''),
        "session_path": session_path,
        "images": images_data
    }


def run_data_collection(collect_time):
    """在单独线程中运行数据采集"""
    global collection_status

    try:
        collector = ADBManager()
        success = collector.collect_data(collect_time)

        if success:
            # 统计采集的图片数量
            latest_session = get_latest_collection_session()
            if latest_session:
                image_count = count_images_in_session(latest_session)
                collection_status["last_collection_count"] = image_count
                collection_status["last_collection_time"] = datetime.now().isoformat()
            else:
                collection_status["last_collection_count"] = 0
        else:
            collection_status["last_collection_count"] = 0

    except Exception as e:
        print(f"Data collection error: {e}")
        collection_status["last_collection_count"] = 0
    finally:
        collection_status["is_collecting"] = False
        collection_status["current_collection_thread"] = None
    print(collection_status)


@router.get("/collect")
async def collect_data(
        time: float = Query(3.0, description="数据采集时间（秒）", ge=0.1, le=60.0)
):
    """
    执行数据采集

    参数:
    - time: 数据采集时间（秒）

    返回:
    - 采集状态和图片数量
    """
    global collection_status

    if collection_status["is_collecting"]:
        raise HTTPException(status_code=400, detail="数据采集中，请等待当前采集完成")

    # 开始数据采集
    collection_status["is_collecting"] = True

    # 在单独线程中运行数据采集
    collection_thread = threading.Thread(
        target=run_data_collection,
        args=(time,)
    )
    collection_thread.daemon = True
    collection_thread.start()

    collection_status["current_collection_thread"] = collection_thread

    return JSONResponse(content={
        "status": "success",
        "message": f"数据采集已开始，预计耗时 {time} 秒",
        "collect_time": time,
        "collection_id": datetime.now().isoformat()
    })


@router.get("/collect/status")
async def get_collection_status():
    """获取数据采集状态"""
    return JSONResponse(content={
        "status": "success",
        "is_collecting": collection_status["is_collecting"],
        "last_collection_time": collection_status["last_collection_time"],
        "last_collection_count": collection_status["last_collection_count"]
    })


@router.get("/image")
async def get_collection_images(
        id: Optional[str] = Query(None, description="数据ID（时间戳排序）")
):
    """
    获取数据采集的图片

    参数:
    - id: 数据采集时间戳ID，如果为空则返回最新采集的数据

    返回:
    - 前6个Camera的图片（base64编码）
    """
    try:
        images_data = get_images_by_id(id)

        if not images_data:
            raise HTTPException(status_code=404, detail="未找到指定的数据采集会话")

        return JSONResponse(content={
            "status": "success",
            "session_timestamp": images_data["session_timestamp"],
            "image_count": len(images_data["images"]),
            "images": images_data["images"]
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取图片失败: {str(e)}")


@router.get("/sessions")
async def get_all_sessions():
    """获取所有数据采集会话"""
    try:
        sessions = get_all_collection_sessions()

        return JSONResponse(content={
            "status": "success",
            "sessions_count": len(sessions),
            "sessions": [
                {
                    "timestamp": session["timestamp"],
                    "session_name": session["session_name"],
                    "image_count": session["image_count"],
                    "creation_time": datetime.fromtimestamp(session["creation_time"]).isoformat()
                }
                for session in sessions
            ]
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取会话列表失败: {str(e)}")


@router.get("/stats")
async def get_collection_stats():
    """获取数据采集统计信息"""
    try:
        sessions = get_all_collection_sessions()
        total_images = sum(session["image_count"] for session in sessions)

        return JSONResponse(content={
            "status": "success",
            "total_sessions": len(sessions),
            "total_images": total_images,
            "data_root": DATA_ROOT,
            "service_data_path": SERVICE_DATA_PATH,
            "collection_status": {
                "is_collecting": collection_status["is_collecting"],
                "last_collection_time": collection_status["last_collection_time"],
                "last_collection_count": collection_status["last_collection_count"]
            }
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计信息失败: {str(e)}")


@router.get("/health")
async def data_health_check():
    """数据采集模块健康检查"""
    data_root_exists = os.path.exists(DATA_ROOT)
    service_path_exists = os.path.exists(SERVICE_DATA_PATH)

    sessions_count = len(get_all_collection_sessions())

    return JSONResponse(content={
        "status": "success",
        "data_root_exists": data_root_exists,
        "service_path_exists": service_path_exists,
        "sessions_count": sessions_count,
        "collection_status": collection_status["is_collecting"]
    })


# 初始化检查
def initialize_data_module():
    """初始化数据模块"""
    print("初始化数据采集模块...")

    # 创建数据目录（如果不存在）
    os.makedirs(DATA_ROOT, exist_ok=True)
    os.makedirs(SERVICE_DATA_PATH, exist_ok=True)

    print(f"数据根目录: {DATA_ROOT}")
    print(f"服务数据路径: {SERVICE_DATA_PATH}")

    # 检查现有会话
    sessions = get_all_collection_sessions()
    print(f"发现 {len(sessions)} 个数据采集会话")

    for session in sessions[:3]:  # 显示最近3个会话
        print(f"  - {session['session_name']}: {session['image_count']} 张图片")


# 在模块加载时初始化
initialize_data_module()
